import { Language, RoleMessageType, RoleType } from "../enum";
import { ChatMessage } from "../message/chatMessage";
import { SystemMessage } from "../message/systemMessage";
import { breakMessages } from "../message/breakMessages";
import { getPromptFromKey, promptTemplateDict } from "../prompt";
import { ChatAgent, RawChatAgentConfig, ChatAgentResponse } from "./chatAgent";

/**
 * 分析代码的Agent
 */
export class CodeAnalyzerAgent extends ChatAgent {
    constructor(config: RawChatAgentConfig) {
        const systemMessage = new SystemMessage({
            roleName: "CodeAnalyzer",
            roleType: RoleType.Assistant,
            roleMessageType: RoleMessageType.AgentSystem,
            content: "你按照需求分析用户提供的代码"
        });
        super({ ...config, systemMessage });
    }

    async step(inputMessage: ChatMessage): Promise<ChatAgentResponse> {
        this.reset();
        const prompt = getPromptFromKey(
            promptTemplateDict.codeTemplate.codeAnalyzer,
            this.language
        ).format({
            code: inputMessage.content
        });
        const codeMessage = new ChatMessage({
            roleName: this.roleName,
            roleType: this.roleType,
            roleMessageType: RoleMessageType.AgentReceiving,
            content: prompt
        });
        return await super.step(codeMessage);
    }
}

interface CodeHandlerAgentTemplateDict {
    /** 主模板 */
    mainTemplate: Record<Language, string>;
    /** 步骤模板 */
    stepTemplate: Record<Language, string>;
    /** 步骤统计模板 */
    finalStatisticalTemplate: Record<Language, string>;
    /** 步骤总结模板 */
    finalInstructionsTemplate: Record<Language, string>;
}

/**
 * 通用代码处理Agent
 */
export class CodeHandlerAgent extends ChatAgent {
    public maxBreak: number;
    public projectName: string;
    public projectDescription: string;

    constructor(
        config: {
            maxBreakPower?: number;
            projectName?: string;
            projectDescription?: string;
        } & RawChatAgentConfig,
        public promptTemplateDict: CodeHandlerAgentTemplateDict // 不同的模板传入
    ) {
        const systemMessage = new SystemMessage({
            roleName: "CodeHandler",
            roleType: RoleType.Assistant,
            roleMessageType: RoleMessageType.AgentSystem,
            content: "你帮助用户处理代码摘要"
        });
        super({ ...config, systemMessage });

        if (config.maxBreakPower && config.maxBreakPower > 4) {
            throw new Error("maxBreak should be less than or equal to 4");
        }
        this.maxBreak = 2 ** (config.maxBreakPower || 2);
        this.projectName = config.projectName || "";
        this.projectDescription = config.projectDescription || "";
    }

    async step(inputMessage: ChatMessage): Promise<ChatAgentResponse> {
        this.reset();
        const content = await this.handelMessages(inputMessage);

        const prompt = getPromptFromKey(
            this.promptTemplateDict.mainTemplate, // 使用传入的模板
            this.language
        ).format({
            code: content
        });
        const codeMessage = new ChatMessage({
            roleName: this.roleName,
            roleType: this.roleType,
            roleMessageType: RoleMessageType.AgentReceiving,
            content: prompt
        });

        if (this.isExceedTokenLimit(codeMessage.content)) {
            throw new Error("Token count exceeds limit");
        }

        this.reset();
        return await super.step(codeMessage);
    }

    private async handelMessages(inputMessage: ChatMessage): Promise<string> {
        const messages = (
            await breakMessages(
                inputMessage,
                this.modelBackend,
                this.maxBreak,
                (message, index) => {
                    const templateKey =
                        index === 0
                            ? this.promptTemplateDict.mainTemplate
                            : this.promptTemplateDict.stepTemplate;

                    return getPromptFromKey(templateKey, this.language).format({
                        codeCount: `${index}`,
                        codeIndex: `${index}`,
                        code: message.content,
                        projectName: this.projectName,
                        projectDescription: this.projectDescription
                    });
                }
            )
        ).map((message) => message.content);

        // 如果只有一个代码，直接返回
        if (messages.length === 1) return messages[0];

        const summaryResponses = [];
        for (let index = 0; index < messages.length; index++) {
            const code = messages[index];
            const p = getPromptFromKey(this.promptTemplateDict.stepTemplate, this.language).format({
                codeCount: `${messages.length}`,
                codeIndex: `${index + 1}`,
                code
            });
            const chatMessage = new ChatMessage({
                roleName: this.roleName,
                roleType: this.roleType,
                roleMessageType: RoleMessageType.AgentReceiving,
                content: p
            });
            super.reset();
            const response = await super.step(chatMessage);
            summaryResponses.push(response);
        }

        const prompt1 = getPromptFromKey(
            this.promptTemplateDict.finalInstructionsTemplate,
            this.language
        ).format({});

        const prompt2Array = [];
        for (let index = 0; index < summaryResponses.length; index++) {
            const content = summaryResponses[index].messages.pop();
            const p = getPromptFromKey(
                this.promptTemplateDict.finalStatisticalTemplate,
                this.language
            ).format({
                codeIndex: `${index + 1}`,
                code: content?.content || ""
            });
            prompt2Array.push(p);
        }
        const prompt2 = prompt2Array.join("\n");
        return prompt1 + prompt2;
    }
}

/**
 * 提供代码摘要的Agent
 */
export class CodeAbstractorAgent extends CodeHandlerAgent {
    constructor(
        config: {
            maxBreakPower?: number;
            projectName?: string;
            projectDescription?: string;
        } & RawChatAgentConfig
    ) {
        super(config, {
            mainTemplate: promptTemplateDict.codeTemplate.codeAbstractor,
            stepTemplate: promptTemplateDict.codeTemplate.codeAbstractorStep,
            finalStatisticalTemplate: promptTemplateDict.codeTemplate.codeAbstractorStepStatistical,
            finalInstructionsTemplate:
                promptTemplateDict.codeTemplate.codeAbstractorStepInstructions
        });
    }
}

/**
 * 文件夹代码摘要的Agent
 */
export class CodeFolderAbstractorAgent extends CodeHandlerAgent {
    constructor(
        config: {
            maxBreakPower?: number;
            projectName?: string;
            projectDescription?: string;
        } & RawChatAgentConfig
    ) {
        super(config, {
            mainTemplate: promptTemplateDict.codeTemplate.codeFolderAbstractor,
            stepTemplate: promptTemplateDict.codeTemplate.codeFolderAbstractorStep,
            finalStatisticalTemplate:
                promptTemplateDict.codeTemplate.codeFolderAbstractorStepStatistical,
            finalInstructionsTemplate:
                promptTemplateDict.codeTemplate.codeFolderAbstractorStepInstructions
        });
    }
}

/**
 * 总结代码的Agent
 */
export class CodeSummarizerAgent extends ChatAgent {
    constructor(config: RawChatAgentConfig) {
        const systemMessage = new SystemMessage({
            roleName: "CodeSummarizer",
            roleType: RoleType.Assistant,
            roleMessageType: RoleMessageType.AgentSystem,
            content: "你按照需求总结用户提供的代码"
        });
        super({ ...config, systemMessage });
    }

    async step(inputMessage: ChatMessage): Promise<ChatAgentResponse> {
        this.reset();
        const prompt = getPromptFromKey(
            promptTemplateDict.codeTemplate.codeSummarizer,
            this.language
        ).format({
            code: inputMessage.content
        });
        const codeMessage = new ChatMessage({
            roleName: this.roleName,
            roleType: this.roleType,
            roleMessageType: RoleMessageType.AgentReceiving,
            content: prompt
        });
        return await super.step(codeMessage);
    }
}

/**
 * 总结代码总结的Agent
 */
export class CodeSummarizerSummaryAgent extends ChatAgent {
    constructor(config: RawChatAgentConfig) {
        const systemMessage = new SystemMessage({
            roleName: "CodeSummarizerSummary",
            roleType: RoleType.Assistant,
            roleMessageType: RoleMessageType.AgentSystem,
            content: "你帮助用户提炼更高层次的代码总结"
        });
        super({ ...config, systemMessage });
    }

    async step(inputMessage: ChatMessage): Promise<ChatAgentResponse> {
        this.reset();
        const prompt = getPromptFromKey(
            promptTemplateDict.codeTemplate.codeSummarizerSummary,
            this.language
        ).format({
            summarys: inputMessage.content
        });
        const codeMessage = new ChatMessage({
            roleName: this.roleName,
            roleType: this.roleType,
            roleMessageType: RoleMessageType.AgentReceiving,
            content: prompt
        });
        return await super.step(codeMessage);
    }
}
